# 官方源码参考
# class Conv1D(nn.Module):
#     """
#     1D-convolutional layer as defined by Radford et al. for OpenAI GPT (and also used in GPT-2).

#     Basically works like a linear layer but the weights are transposed.

#     Args:
#         nf (`int`): The number of output features.
#         nx (`int`): The number of input features.
#     """

#     def __init__(self, nf, nx):
#         super().__init__()
#         self.nf = nf
#         self.nx = nx
#         self.weight = nn.Parameter(torch.empty(nx, nf))
#         self.bias = nn.Parameter(torch.zeros(nf))
#         nn.init.normal_(self.weight, std=0.02)

#     def __repr__(self) -> str:
#         return "Conv1D(nf={nf}, nx={nx})".format(**self.__dict__)

#     def forward(self, x):
#         size_out = x.size()[:-1] + (self.nf,)
#         x = torch.addmm(self.bias, x.view(-1, x.size(-1)), self.weight)
#         x = x.view(size_out)
#         return x



import torch
import torch.nn as nn

# 输入：batch_size=2, input_dim=5
x = torch.tensor([[1., 2., 3., 4., 5.],
                  [5., 4., 3., 2., 1.]])
print(f"x.shape: {x.shape}")  # torch.Size([2, 5])




class Conv1D(nn.Module):
    def __init__(self, nf, nx):
        super().__init__()
        self.nf = nf  # output dim
        self.nx = nx  # input dim
        self.weight = nn.Parameter(torch.arange(nx * nf).reshape(nx, nf).float())
        self.bias = nn.Parameter(torch.ones(nf))  # 设为1方便观察

    def forward(self, x):
        size_out = x.size()[:-1] + (self.nf,)
        x = x.view(-1, x.size(-1))                     # 变成二维 [batch, input]
        x = torch.addmm(self.bias, x, self.weight)     # 注意这里是 x @ W
        x = x.view(size_out)
        return x

conv = Conv1D(nf=3, nx=5)
out_conv = conv(x)
print("Conv1D output:")
print(out_conv)


